package scu.maqiang.optimalcontrol;

import scu.maqiang.fes.BVPType;
import scu.maqiang.fes.FES2T31;
import scu.maqiang.mesh.*;
import scu.maqiang.numeric.*;

import java.util.Arrays;
import java.util.function.Function;

/**
 * 最优控制问题之最优稳态热源
 */
public class OptimalHeatSource {
    public static void main(String[] args) {
        Mesh2T3 mesh = new Mesh2T3();
        mesh.square2D(1024, 1024, T3Type.Left);
        FES2T31 fs = new FES2T31(mesh);
        int dof = fs.GetNdof();
        double[] y = new double[dof];   //状态变量
        double[] u = new double[dof];   //控制变量
        double[] p = new double[dof];   //伴随变量

        double[] uprev = new double[dof];  //记录上次的控制变量
        Arrays.fill(uprev, 1.0);

        double alpha = 1.0;
        double[] ngradF = new double[dof];
        ScalarFunc desiredFunc = (xy, label, param) -> Math.sin(10 * Math.PI * xy[0]) * Math.sin(10 * Math.PI * xy[1]);
        double[] z = mesh.functionInMesh(desiredFunc, null);  //设置需要的状态量

        SRMatrix lap = new SRMatrix(dof);
        fs.assembleStiff(new double[]{1.0}, BVPType.COMMON, lap);
        fs.applyBC_MBN(lap, Direct.X, 1, 2, 3, 4);
        DirectSSolver solver = new DirectSSolver(lap);
        solver.Factorize_Cholesky(1);

        double[] RHS = new double[dof];
        double eps = 1.0e-6;

        //定义目标泛函
        CostFunctional Linear_Quadratic_Functional = (state, control) -> {
            double[] zz = state.clone();
            Arrays.setAll(zz, i -> state[i] - z[i]);
            double temp1 = fs.computeL2Norm(zz);
            double temp2 = fs.computeL2Norm(control);
            return 0.5 * temp1 * temp1 + 0.5 * alpha * temp2 * temp2;
        };

        //定义由控制变量计算状态变量的方法
        Operator controlToState = (control, state) -> {
            Arrays.fill(RHS, 0.0);
            fs.assembleSource(new double[][]{control}, BVPType.COMMON, RHS);
            fs.applyBC_MBN(RHS, Direct.X, 0.0, 1, 2, 3, 4);
            solver.Solve_Cholesky(RHS, state, 0);
            //solver.PCGSSOR(controlRHS, state, 1.5, 1);
        };

        //定义由状态变量计算伴随变量的方法
        Operator stateToAdjoint = (state, adjoint) -> {
            Arrays.fill(RHS, 0.0);
            double[] vecRHS = RHS.clone();
            Arrays.setAll(vecRHS, i -> state[i] - z[i]);
            //System.out.println("min: " + MVO.min(vecRHS) + "\t max: " + MVO.max(vecRHS));
            fs.assembleSource(new double[][]{vecRHS}, BVPType.COMMON, RHS);
            fs.applyBC_MBN(RHS, Direct.X, 0.0, 1, 2, 3, 4);
            solver.Solve_Cholesky(RHS, adjoint, 0);
        };

        Function<double[], Double> J = control -> {
            controlToState.operate(control, y);
            double[] zz = y.clone();
            Arrays.setAll(zz, i -> y[i] - z[i]);
            double temp1 = fs.computeL2Norm(zz);
            double temp2 = fs.computeL2Norm(control);
            return 0.5 * temp1 * temp1 + 0.5 * alpha * temp2 * temp2;
        };

        Function<Double[], Double[]> DJ = control -> {
            stateToAdjoint.operate(y, p);
            Arrays.setAll(ngradF, i -> alpha * u[i] + p[i]);
            return null;
        };

        double time = 0.0;
        mesh.toTecplot("State.dat", time, y);
        while (MVO.distance(u, uprev) > eps) {

            Arrays.setAll(uprev, i -> u[i]);
            //1.由控制计算状态变量
            controlToState.operate(u, y);

            //2. 由状态计算伴随变量
            stateToAdjoint.operate(y, p);

            //3. 计算下降方向, 这里使用负梯度方向
            Arrays.setAll(ngradF, i -> -(alpha * u[i] + p[i]));
            double decentNorm2 = fs.integrate(ngradF);
            decentNorm2 *= decentNorm2;

            //4. 使用Armijo准则计算步长
            double step = OptimalControl.ArmijoStep(Linear_Quadratic_Functional, controlToState, ngradF, decentNorm2,
                    0.5, 1.0e-4, y, u);

            //5. 更新控制
            Arrays.setAll(u, i -> u[i] + step * ngradF[i]);

            time = time + 1.0;
            System.out.println("min: " + MVO.min(y) + "\t max: " + MVO.max(y));
            mesh.toTecplot("State.dat", time, y);
        }
        System.out.println("min: " + MVO.min(y) + "\t max: " + MVO.max(y));
        System.out.println("min: " + MVO.min(u) + "\t max: " + MVO.max(u));
        mesh.toTecplot("FinalState.dat", y);
        //OptimalHeatSourceV2();
    }

    public static void OptimalHeatSourceV2() {
        Mesh2T3 mesh = new Mesh2T3();
        mesh.square2D(100, 100, T3Type.Left);
        FES2T31 fs = new FES2T31(mesh);
        int dof = fs.GetNdof();
        double[] y = new double[dof];   //状态变量
        double[] u = new double[dof];   //控制变量
        double[] p = new double[dof];   //伴随变量

        double[] uprev = new double[dof];  //记录上次的控制变量
        Arrays.fill(uprev, 1.0);

        double alpha = 1.0;
        double[] ngradF = new double[dof];
        ScalarFunc desiredFunc = (xy, label, param) -> Math.sin(2 * Math.PI * xy[0]) * Math.sin(2 * Math.PI * xy[1]);
        double[] z = mesh.functionInMesh(desiredFunc, null);  //设置需要的状态量

        SRMatrix lap = new SRMatrix(dof);
        fs.assembleStiff(new double[]{1.0}, BVPType.COMMON, lap);
        fs.applyBC_MBN(lap, Direct.X, 1, 2, 3, 4);
        DirectSSolver solver = new DirectSSolver(lap);
        solver.Factorize_Cholesky(1);

        double[] RHS = new double[dof];
        double eps = 1.0e-6;

        //定义目标泛函
        CostFunctional Linear_Quadratic_Functional = (state, control) -> {
            double[] zz = state.clone();
            Arrays.setAll(zz, i -> state[i] - z[i]);
            double temp1 = fs.computeL2Norm(zz);
            double temp2 = fs.computeL2Norm(control);
            return 0.5 * temp1 * temp1 + 0.5 * alpha * temp2 * temp2;
        };

        //定义由控制变量计算状态变量的方法
        Operator controlToState = (control, state) -> {
            Arrays.fill(RHS, 0.0);
            fs.assembleSource(new double[][]{control}, BVPType.COMMON, RHS);
            fs.applyBC_MBN(RHS, Direct.X, 0.0, 1, 2, 3, 4);
            solver.Solve_Cholesky(RHS, state, 0);
            //solver.PCGSSOR(controlRHS, state, 1.5, 1);
        };

        //定义由状态变量计算伴随变量的方法
        Operator stateToAdjoint = (state, adjoint) -> {
            Arrays.fill(RHS, 0.0);
            double[] vecRHS = RHS.clone();
            Arrays.setAll(vecRHS, i -> state[i] - z[i]);
            //System.out.println("min: " + MVO.min(vecRHS) + "\t max: " + MVO.max(vecRHS));
            fs.assembleSource(new double[][]{vecRHS}, BVPType.COMMON, RHS);
            fs.applyBC_MBN(RHS, Direct.X, 0.0, 1, 2, 3, 4);
            solver.Solve_Cholesky(RHS, adjoint, 0);
        };

        ScalarFunc J = (control, region, param) -> {
            controlToState.operate(control, y);
            double[] zz = y.clone();
            Arrays.setAll(zz, i -> y[i] - z[i]);
            double temp1 = fs.computeL2Norm(zz);
            double temp2 = fs.computeL2Norm(control);
            return 0.5 * temp1 * temp1 + 0.5 * alpha * temp2 * temp2;
        };

        VectorFunc DJ = (control, region, param) -> {
            stateToAdjoint.operate(y, p);
            Arrays.setAll(ngradF, i -> alpha * u[i] + p[i]);
            return ngradF;
        };

        controlToState.operate(u, y);
        stateToAdjoint.operate(y, p);
        double time = 0.0;
        mesh.toTecplot("State.dat", time, y);

        while (MVO.distance(u, uprev) > eps) {

            Arrays.setAll(uprev, i -> u[i]);
//            //1.由控制计算状态变量
//
//            //2. 由状态计算伴随变量
//
//            //3. 计算下降方向, 这里使用负梯度方向
            Arrays.setAll(ngradF, i -> -(alpha * u[i] + p[i]));
//            double decentNorm2 = fs.integrate(ngradF);
//            decentNorm2 *= decentNorm2;
//
//            //4. 使用Armijo准则计算步长
//            double step = OptimalControl.ArmijoStep(Linear_Quadratic_Functional, controlToState, ngradF, decentNorm2,
//                    0.5, 1.0e-4, y, u);

            double step = NLOPT.ArmijoStep(J, DJ, u, ngradF, null, null);

            //5. 更新控制
            Arrays.setAll(u, i -> u[i] + step * ngradF[i]);

            time = time + 1.0;
            mesh.toTecplot("State.dat", time, y);
        }
        System.out.println("min: " + MVO.min(y) + "\t max: " + MVO.max(y));
        System.out.println("min: " + MVO.min(u) + "\t max: " + MVO.max(u));
        mesh.toTecplot("FinalState.dat", y);
    }

}
